import os

import numpy as np
import torch
import pickle


class ReplayBuffer:
    def __init__(self, state_dim, action_dim, device='cpu', max_size=int(1e6)):
        self.max_size = max_size
        self.ptr = 0
        self.size = 0

        self.state = np.zeros((max_size, state_dim))
        self.action = np.zeros((max_size, action_dim))
        self.next_state = np.zeros((max_size, state_dim))
        self.reward = np.zeros((max_size, 1))
        self.done = np.zeros((max_size, 1))

        self.device = device

    def push(self, state, action, next_state, reward, done):
        self.state[self.ptr] = state
        self.action[self.ptr] = action
        self.next_state[self.ptr] = next_state
        self.reward[self.ptr] = reward
        self.done[self.ptr] = done
        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def sample(self, batch_size):
        ind = np.random.randint(0, self.size, size=batch_size)
        return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.done[ind]).to(self.device)
        )

    def save(self, filename, save_dir='./buffers'):
        data = {
            'state': self.state[:self.size],
            'action': self.action[:self.size],
            'next_state': self.next_state[:self.size],
            'reward': self.reward[:self.size],
            'done': self.done[:self.size],
        }

        if not os.path.exists(save_dir):
            os.mkdir(save_dir)

        with open(f'./{save_dir}/{filename}.pkl', 'wb') as f:
            pickle.dump(data, f)

    def load(self, filename, size=-1, save_dir='./buffers'):
        with open(f'./{save_dir}/{filename}.pkl', 'rb') as f:
            data = pickle.load(f)

        # Adjust crt_size if we're using a custom size
        size = min(int(size), self.max_size) if size > 0 else self.max_size
        self.size = min(data["reward"].shape[0], size)

        self.state[:self.size] = data["state"][:self.size]
        self.action[:self.size] = data["action"][:self.size]
        self.next_state[:self.size] = data["next_state"][:self.size]
        self.reward[:self.size] = data["reward"][:self.size]
        self.done[:self.size] = data["done"][:self.size]

    def normalize_states(self, eps=1e-3):
        mean = self.state[:self.size].mean(0, keepdims=True)
        std = self.state[:self.size].std(0, keepdims=True) + eps
        self.state[:self.size] = (self.state[:self.size] - mean) / std
        self.next_state[:self.size] = (self.next_state[:self.size] - mean) / std
        return mean, std
